from typing import Sequence
import functools

import quimb as qu
from quimb.tensor.tensor_core import (
	Tensor,
	rand_uuid, oset, tags_to_oset,
)
from quimb.tensor import tensor_network_1d_compress
from quimb.tensor.tensor_1d import TensorNetwork1DOperator, TensorNetwork1DFlat, MatrixProductOperator

from autoray import do

__all__ = ["PTM_diag_MPO", "PTM_SPLM_MPO", "PTM_IDM_MPO"]

_I_o4 = qu.eye(4).reshape(1, 1, 4, 4)

_G0 = do("diag", [1.0, 0, 0, 0], like=None)
_G1 = do("diag", [0.0, 1, 0, 0], like=None)
_G2 = do("diag", [0.0, 0, 1, 0], like=None)
_G3 = do("diag", [0.0, 0, 0, 1], like=None)

_G_o3 = qu.qu([_G0, _G1, _G2, _G3])

class PTM_diag_MPO(TensorNetwork1DOperator, TensorNetwork1DFlat):
	"""Initialise a TensorNetwork of Sparse Pauli Lindblad noise model(SPLM)
	params_1qu: dict
		A dictionary containing the parameters of noise model. The key that indicates
		the location of the qubit site should be int, and the corresponding value are arrays-like
		parameters.
	params_2qu: dict
		A dictionary containing the parameters of noise model. The key that indicates
		the location of the two qubit site should be tuple, and the corresponding value are arrays-like
		parameters.
	L: int, optional
        The number of sites the MPO should be defined on.
    tags : str or sequence of str, optional
        Global tags to attach to all tensors.
	upper_ind_id: str
        A string specifying how to label the upper physical site indices.
        Should contain a ``'{}'`` placeholder. It is used to generate the
        actual indices like: ``map(upper_ind_id.format, range(len(arrays)))``.
	lower_ind_id: str
        A string specifying how to label the lower physical site indices.
        Should contain a ``'{}'`` placeholder. It is used to generate the
        actual indices like: ``map(lower_ind_id.format, range(len(arrays)))``.
	site_tag_id: str
        A string specifying how to tag the tensors at each site. Should
        contain a ``'{}'`` placeholder. It is used to generate the actual tags
        like: ``map(site_tag_id.format, range(len(arrays)))``.
	"""
	@classmethod
	def from_param_dict(
			cls,
			params_1qu: dict[int, ...],
			params_2qu: dict[Sequence[int], ...],
			*,
			L: int = None,
			tags=None,
			upper_ind_id: str = "k{}",
			lower_ind_id: str = "b{}",
			site_tag_id: str = "I{}",
			**tn_opts
	):
		if not isinstance(params_1qu, dict) and not isinstance(params_2qu, dict):
			raise ValueError
		tn = cls.new(L=L,
		             upper_ind_id=upper_ind_id,
		             lower_ind_id=lower_ind_id,
		             site_tag_id=site_tag_id,
		             **tn_opts)
		tags = tags_to_oset(tags)

		ptm_mpo_2qu_ = functools.partial(cls.ptm_mpo_2qu,
		                                 upper_ind_id=upper_ind_id,
		                                 lower_ind_id=lower_ind_id,
		                                 site_tag_id=site_tag_id)
		for site, param_1qu in params_1qu.items():
			t = Tensor(cls.ptm_mat_1qu(param_1qu), inds=(upper_ind_id.format(site), lower_ind_id.format(site)),
			           tags=tags | oset([site_tag_id.format(site)]))
			L = max(L, site)
			tn.add_tensor(t)
		for site, param_2qu in params_2qu.items():
			i, j = site[0], site[1]
			t = ptm_mpo_2qu_(param_2qu, i, j, tags=tags)
			L = max(L, i, j)
			tn.add_tensor_network(t)
		tensor_network_1d_compress(tn, method='direct', inplace=True)
		return tn
	@classmethod
	def from_adj_param_mat(
			cls,
			param_mat_1qu,
			param_mat_2qu,
			*,
			L: int = None,
			tags=None,
			upper_ind_id: str = "k{}",
			lower_ind_id: str = "b{}",
			site_tag_id: str = "I{}",
			like=None,
			**tn_opts
	):
		"""Initialise a TensorNetwork of Sparse Pauli Lindblad noise model(SPLM) where the
		two-qubit SPLM parameters are adjacent.
		param_mat_1qu: array-like
				An array containing the all parameters of noise model, where the i-th row indicates
				the one-qubit noise parameters of i-th qubit,
		param_mat_2qu: array-like
				An array containing the all parameters of noise model, where the i-th row indicates
				the two-qubit noise parameters of i-th qubit and (i+1)-th qubit,
		L: int, optional
		        The number of sites the MPO should be defined on.
		upper_ind_id: str
		lower_ind_id: str
		site_tag_id: str
		"""
		if L is None:
			assert (L := param_mat_1qu.shape[0]) == param_mat_2qu.shape[0] + 1
		else:
			assert L == param_mat_1qu.shape[0] == param_mat_2qu.shape[0] + 1

		tn = cls.new(L=L,
		             upper_ind_id=upper_ind_id,
		             lower_ind_id=lower_ind_id,
		             site_tag_id=site_tag_id,
		             **tn_opts)
		tags = tags_to_oset(tags)

		sites = tuple(range(L))
		upper_inds, lower_inds, site_tags = \
			map(upper_ind_id.format, sites), map(lower_ind_id.format, sites), map(site_tag_id.format, sites)

		ptm_mat_1qu = cls.ptm_mat_1qu
		ptm_mat_2qu = cls.ptm_mat_2qu

		k1_ix, n_ix = rand_uuid(), rand_uuid()
		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
		t1 = Tensor(ptm_mat_1qu(param_mat_1qu[0], like=like),
		            inds=(k1_ix, lower_ind), tags=tags | oset([site_tag]))
		t2 = Tensor(ptm_mat_2qu(param_mat_2qu[0], like=like),
		            inds=(n_ix, upper_ind, k1_ix), tags=tags | oset([site_tag]))
		tn.add_tensor(t1 @ t2)
		p_ix = n_ix
		for i in range(1, L - 1):
			k1_ix, k2_ix, n_ix = rand_uuid(), rand_uuid(), rand_uuid()
			upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
			t1 = Tensor(ptm_mat_1qu(param_mat_1qu[i], like=like),
			            inds=(k1_ix, lower_ind), tags=tags | oset([site_tag]))
			if i % 2:
				t2 = Tensor(_G_o3,
				            inds=(p_ix, k2_ix, k1_ix), tags=tags | oset([site_tag]))
				t3 = Tensor(ptm_mat_2qu(param_mat_2qu[i], like=like),
				            inds=(n_ix, upper_ind, k2_ix), tags=tags | oset([site_tag]))
			else:
				t2 = Tensor(ptm_mat_2qu(param_mat_2qu[i], like=like),
				            inds=(n_ix, k2_ix, k1_ix), tags=tags | oset([site_tag]))
				t3 = Tensor(_G_o3,
				            inds=(p_ix, upper_ind, k2_ix), tags=tags | oset([site_tag]))
			tn.add_tensor(t1 @ t2 @ t3)
			p_ix = n_ix
		k1_ix = rand_uuid()
		upper_ind, lower_ind, site_tag = next(upper_inds), next(lower_inds), next(site_tags)
		t1 = Tensor(ptm_mat_1qu(param_mat_1qu[-1], like=like),
		            inds=(k1_ix, lower_ind), tags=tags | oset([site_tag]))
		t2 = Tensor(_G_o3,
		            inds=(p_ix, upper_ind, k1_ix), tags=tags | oset([site_tag]))
		tn.add_tensor(t1 @ t2)
		return tn

	@staticmethod
	def ptm_mat_1qu(param, like=None):
		raise NotImplementedError
	@staticmethod
	def ptm_mat_2qu(param, like=None):
		raise NotImplementedError
	@staticmethod
	def ptm_mpo_2qu(
			param, i, j,
			tags=None,
			upper_ind_id="k{}",
			lower_ind_id="b{}",
			site_tag_id="I{}",
			**tn_opts
	) -> MatrixProductOperator:
		i, j = min(i, j), max(i, j)
		def gen_arrays():
			yield PTM_diag_MPO.ptm_mat_2qu(param)
			for _ in range(i + 1, j):
				yield _I_o4
			yield _G_o3
		return MatrixProductOperator(gen_arrays(),
		                             sites=tuple(range(i, j + 1)),
		                             tags=tags,
		                             upper_ind_id=upper_ind_id,
		                             lower_ind_id=lower_ind_id,
		                             site_tag_id=site_tag_id,
		                             **tn_opts)

class PTM_SPLM_MPO(PTM_diag_MPO):
	"""Initialise a TensorNetwork of Sparse Pauli Lindblad noise model(SPLM)
	lams_1qu: dict
		A dictionary containing the parameters of one-qubit SPLM. The key that indicates
		the location of the qubit site should be int, and the corresponding arrays-like
		parameters should be of length 3 (X, Y, Z, respectively), e.g.,
		``lams = {
			# qubit 0, with parameters lamX = 0.01, lamY = 0.003, lamZ = 0
			0 : np.array([0.01, 0.003, 0]),
			# qubit 3, with parameters lamX = 0, lamY = 0.001, lamZ = 0
			3 : [0, 0.001, 0],
		}``
	lams_1qu: dict
		A dictionary containing the parameters of two-qubit SPLM. The key that indicates
		the location of the two qubit site should be tuple, and the corresponding arrays-like
		parameters should be of length 9 (XX, XY, XZ, YX, YY, YZ, ZX, ZY, ZZ, respectively), e.g.,
		``lams = {
			# qubit 0 and qubit 2, with parameters
			# lamXX = 0.01, lamXY = lamXZ = 0
			# lamYX = 0, lamYY = 0.002, lamYZ = 0
			# lamZX = 0.005, lamZY = 0, lamZZ = 0
			(0, 2) : [0.01, 0, 0, 0, 0.002, 0, 0005, 0, 0],
		}``
	L: int, optional
        If not given, this is taken as the max key value of ``lams``.
	"""
	@staticmethod
	def ptm_mat_1qu(lams, like=None):
		lam1, lam2, lam3 = lams
		# zeta1 = jnp.exp(-2.0 * (lam2 + lam3))
		# zeta2 = jnp.exp(-2.0 * (lam1 + lam3))
		# zeta3 = jnp.exp(-2.0 * (lam1 + lam2))
		# return jnp.diag(jnp.array([1.0, zeta1, zeta2, zeta3]))
		zeta1 = do("exp", -2.0 * (lam2 + lam3), like=like)
		zeta2 = do("exp", -2.0 * (lam1 + lam3), like=like)
		zeta3 = do("exp", -2.0 * (lam1 + lam2), like=like)
		d = do("array", [1.0, zeta1, zeta2, zeta3], like=like)
		return do("diag", d, like=like)
	@staticmethod
	def ptm_mat_2qu(lams, like=None):
		lam11, lam12, lam13, lam21, lam22, lam23, lam31, lam32, lam33 = lams
		# zeta00 = 1.0
		# zeta01 = jnp.exp(-2.0 * (lam12 + lam13 + lam22 + lam23 + lam32 + lam33))
		# zeta02 = jnp.exp(-2.0 * (lam11 + lam13 + lam21 + lam23 + lam31 + lam33))
		# zeta03 = jnp.exp(-2.0 * (lam11 + lam12 + lam21 + lam22 + lam31 + lam32))
		# zeta10 = jnp.exp(-2.0 * (lam21 + lam22 + lam23 + lam31 + lam32 + lam33))
		# zeta11 = jnp.exp(-2.0 * (lam12 + lam13 + lam21 + lam31))
		# zeta12 = jnp.exp(-2.0 * (lam11 + lam13 + lam22 + lam32))
		# zeta13 = jnp.exp(-2.0 * (lam11 + lam12 + lam23 + lam33))
		# zeta20 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam31 + lam32 + lam33))
		# zeta21 = jnp.exp(-2.0 * (lam11 + lam22 + lam23 + lam31))
		# zeta22 = jnp.exp(-2.0 * (lam12 + lam21 + lam23 + lam32))
		# zeta23 = jnp.exp(-2.0 * (lam13 + lam21 + lam22 + lam33))
		# zeta30 = jnp.exp(-2.0 * (lam11 + lam12 + lam13 + lam21 + lam22 + lam23))
		# zeta31 = jnp.exp(-2.0 * (lam11 + lam21 + lam32 + lam33))
		# zeta32 = jnp.exp(-2.0 * (lam12 + lam22 + lam31 + lam33))
		# zeta33 = jnp.exp(-2.0 * (lam13 + lam23 + lam31 + lam32))
		# F0 = jnp.diag(jnp.array([zeta00, zeta10, zeta20, zeta30]))
		# F1 = jnp.diag(jnp.array([zeta01, zeta11, zeta21, zeta31]))
		# F2 = jnp.diag(jnp.array([zeta02, zeta12, zeta22, zeta32]))
		# F3 = jnp.diag(jnp.array([zeta03, zeta13, zeta23, zeta33]))
		# return jnp.array([F0, F1, F2, F3])
		zeta00 = 1.0
		zeta01 = do("exp", -2.0 * (lam12 + lam13 + lam22 + lam23 + lam32 + lam33), like=like)
		zeta02 = do("exp", -2.0 * (lam11 + lam13 + lam21 + lam23 + lam31 + lam33), like=like)
		zeta03 = do("exp", -2.0 * (lam11 + lam12 + lam21 + lam22 + lam31 + lam32), like=like)
		zeta10 = do("exp", -2.0 * (lam21 + lam22 + lam23 + lam31 + lam32 + lam33), like=like)
		zeta11 = do("exp", -2.0 * (lam12 + lam13 + lam21 + lam31), like=like)
		zeta12 = do("exp", -2.0 * (lam11 + lam13 + lam22 + lam32), like=like)
		zeta13 = do("exp", -2.0 * (lam11 + lam12 + lam23 + lam33), like=like)
		zeta20 = do("exp", -2.0 * (lam11 + lam12 + lam13 + lam31 + lam32 + lam33), like=like)
		zeta21 = do("exp", -2.0 * (lam11 + lam22 + lam23 + lam31), like=like)
		zeta22 = do("exp", -2.0 * (lam12 + lam21 + lam23 + lam32), like=like)
		zeta23 = do("exp", -2.0 * (lam13 + lam21 + lam22 + lam33), like=like)
		zeta30 = do("exp", -2.0 * (lam11 + lam12 + lam13 + lam21 + lam22 + lam23), like=like)
		zeta31 = do("exp", -2.0 * (lam11 + lam21 + lam32 + lam33), like=like)
		zeta32 = do("exp", -2.0 * (lam12 + lam22 + lam31 + lam33), like=like)
		zeta33 = do("exp", -2.0 * (lam13 + lam23 + lam31 + lam32), like=like)
		d0 = do("array", [zeta00, zeta10, zeta20, zeta30], like=like)
		d1 = do("array", [zeta01, zeta11, zeta21, zeta31], like=like)
		d2 = do("array", [zeta02, zeta12, zeta22, zeta32], like=like)
		d3 = do("array", [zeta03, zeta13, zeta23, zeta33], like=like)
		F0 = do("diag", d0, like=like)
		F1 = do("diag", d1, like=like)
		F2 = do("diag", d2, like=like)
		F3 = do("diag", d3, like=like)
		return do("array", [F0, F1, F2, F3], like=like)

class PTM_IDM_MPO(PTM_diag_MPO):
	"""Initialise a TensorNetwork of Incoherent Depolarizing noise model(IDM) where the
	two-qubit IDM parameters are adjacent.
	p_vec_1qu: array-like
		An array containing the all parameters of one-qubit IDM, where the i-th row indicates
		the one-qubit IDM parameters of i-th qubit, e.g.,
		``p_vec_1qu = np.array(
			[0.01, 0, 0.001, 0], # qubit 0-3
		)``
	p_vec_2qu: array-like
		An array containing the all parameters of two-qubit IDM, where the i-th row indicates
		the two-qubit IDM parameters of i-th qubit and (i+1)-th qubit, e.g.,
		``p_vec_2qu = np.array(
			[0, 0.001, 0.005], # qubit (0, 1), (1, 2), (2, 3)
		)``
		the number of rows plus 1 should be the same as that of ``lam_mat_1qu``.
	L: int, optional
		If not given, this is taken as the length of ``p_vec_1qu``.
	upper_ind_id: str
	lower_ind_id: str
	site_tag_id: str
	"""
	@staticmethod
	def ptm_mat_1qu(p, like=None):
		rp = 1 - p
		return do("diag", do("array", [1.0, rp, rp, rp], like=like), like=like)
	@staticmethod
	def ptm_mat_2qu(p, like=None):
		rp = 1 - p
		d = do("array", [1, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp, rp], like=like)
		return do("diag", d, like=like)
